from flask import Flask, jsonify, request, send_from_directory
from flask_cors import CORS
import os
import json
import signal
import sys
from gradio_client import Client, file
import subprocess
from subprocess import Popen
import platform

# 设置静态文件目录为public
app = Flask(__name__, static_url_path='', static_folder='public')
CORS(app)

# Global variables
MODELS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
CONFIG_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.json")
PUBLIC_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "public")
WAV_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "wav")
current_model = None
synthesis_process = None

# Ensure directories exist
os.makedirs(MODELS_DIR, exist_ok=True)
os.makedirs(PUBLIC_DIR, exist_ok=True)
os.makedirs(WAV_DIR, exist_ok=True)

def normalize_path(path):
    """统一路径格式，确保使用正确的路径分隔符"""
    if platform.system() == 'Windows':
        return path.replace('/', '\\')
    return path.replace('\\', '/')

def load_config():
    if os.path.exists(CONFIG_FILE):
        with open(CONFIG_FILE, 'r') as f:
            return json.load(f)
    return {"current_model": None}

def save_config(config):
    with open(CONFIG_FILE, 'w') as f:
        json.dump(config, f)

@app.route('/')
def index():
    return app.send_static_file('index.html')

@app.route('/static/<path:path>')
def serve_static(path):
    return send_from_directory(PUBLIC_DIR, path)

@app.route('/wav/<path:path>')
def serve_wav_static(path):
    return send_from_directory(WAV_DIR, path)



@app.route('/api/synthesis/start', methods=['POST'])
def start_synthesis():
    # if not current_model:
    #     return jsonify({"error": "No model selected"}), 400
    
    data = request.get_json()
    text = data.get("text")
    ref_audio_path = data.get("ref_audio_path")
    # if not text:
    #     return jsonify({"error": "No text provided"}), 400
    
    client = Client("http://localhost:9872/")

    result = client.predict(
		text=text,
		text_lang="中文",
        ref_audio_path=file(ref_audio_path),
		aux_ref_audio_paths=[],
		prompt_text="",
		prompt_lang="中文",
		top_k=5,
		top_p=1,
		temperature=1,
		text_split_method="凑四句一切",
		batch_size=20,
		speed_factor=1,
		ref_text_free=False,
		split_bucket=True,
		fragment_interval=0.3,
		seed=-1,
		keep_random=True,
		parallel_infer=True,
		repetition_penalty=1.35,
		api_name="/inference"
    )
    print(result)
    wav_path, identifier = result
    print(wav_path, identifier)
    
    # Create wav directory if it doesn't exist
    wav_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "wav")
    os.makedirs(wav_dir, exist_ok=True)
    
    # Copy the generated wav file to our wav directory
    import shutil
    filename = f"{identifier}.wav"
    target_path = os.path.join(wav_dir, filename)
    try:
        shutil.copy2(wav_path, target_path)
        print(f"Copied WAV file to: {target_path}")
    except Exception as e:
        print(f"Error copying WAV file: {e}")
        return jsonify({"error": f"Failed to copy WAV file: {str(e)}"}), 500
    
    return jsonify({
        "message": "TTS synthesis completed",
        "wav_path": target_path,
        "identifier": identifier
    })

@app.route('/api/synthesis/stop', methods=['POST'])
def stop_synthesis():
    global synthesis_process
    if synthesis_process:
        # Here you would implement the actual stopping mechanism
        synthesis_process = None
        return jsonify({"message": "Synthesis stopped"})
    return jsonify({"message": "No synthesis process running"})

base_gpt_sovits = "D:\\Tool\\GPT-SoVITS-v3lora-20250228\\GPT-SoVITS-v3lora-20250228\\"
python_exec = "runtime\\python.exe"
language = "zh_CN"
p_tts_inference = None
synthesis_process = None
PID_FILE = "synthesis_process.pid"

@app.route('/api/tts/start', methods=['GET'])
def start_tts():
    global synthesis_process
    if synthesis_process:
        synthesis_process.terminate()
        synthesis_process.wait()
        if os.path.exists(PID_FILE):
            os.remove(PID_FILE)

    cmd = f"{base_gpt_sovits}{python_exec} \"{base_gpt_sovits}GPT_SoVITS\\inference_webui_fast.py\" {language}"
    print(cmd)
    synthesis_process = Popen(
        cmd,
        cwd=base_gpt_sovits,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        universal_newlines=True,
        bufsize=1
    )
    
    # Save PID to file
    with open(PID_FILE, 'w') as f:
        f.write(str(synthesis_process.pid))
    
    # Start a thread to monitor the output
    def monitor_output():
        while True:
            line = synthesis_process.stdout.readline()
            if not line:
                break
            print(f"Process output: {line.strip()}")
            # Look for the audio file output pattern
            if ".wav" in line:
                print(f"Found audio: {line}")
                try:
                    # Parse the tuple from the output
                    import ast
                    audio_info = ast.literal_eval(line.strip())
                    if isinstance(audio_info, tuple) and len(audio_info) == 2:
                        audio_path, identifier = audio_info
                        print(f"Found audio: {audio_path}, ID: {identifier}")
                        # You could save this information or trigger some callback here
                except:
                    pass

    import threading
    monitor_thread = threading.Thread(target=monitor_output, daemon=True)
    monitor_thread.start()
    
    return jsonify({"message": cmd, "pid": synthesis_process.pid})

@app.route('/api/tts/stop', methods=['GET'])
def stop_tts():
    global synthesis_process
    try:
        # Try to read PID from file first
        if os.path.exists(PID_FILE):
            with open(PID_FILE, 'r') as f:
                pid = int(f.read().strip())
            try:
                os.kill(pid, signal.SIGTERM)
                # Wait a moment for the process to terminate
                import time
                time.sleep(0.5)
            except ProcessLookupError:
                pass  # Process already terminated
            os.remove(PID_FILE)
            
        # Also handle the process if it exists in memory
        if synthesis_process:
            synthesis_process.terminate()
            synthesis_process.wait()
            synthesis_process = None

        return jsonify({"message": "Process stopped successfully"})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/init_model', methods=['POST'])
def init_model():
    data = request.get_json()
    sovits_path = data.get("pth")
    weights_path = data.get("ckpt")
    if not weights_path:
        return jsonify({"error": "No weights path provided"}), 400

    if not sovits_path:
        return jsonify({"error": "No sovits path provided"}), 400
    
    client = Client("http://localhost:9872/")

    weights_path_result = client.predict(
        weights_path=weights_path,
        api_name="/init_t2s_weights"
    )

    sovits_path_result = client.predict(
        sovits_path=sovits_path,
        prompt_language="中文",
        text_language="中文",
        api_name="/change_sovits_weights"
    )
    return jsonify({"message": "Model initialized successfully", "weights_path": weights_path, "sovits_path": sovits_path, "weights_path_result": weights_path_result, "sovits_path_result": sovits_path_result})

@app.route('/init_pth', methods=['POST'])
def init_pth():
    data = request.get_json()
    sovits_path = data.get("path")
    if not sovits_path:
        return jsonify({"error": "No sovits_path provided"}), 400
   
    client = Client("http://localhost:9872/")
    result = client.predict(
        sovits_path=sovits_path,
        prompt_language="中文",
        text_language="中文",
        api_name="/change_sovits_weights"
    )
    
    print(result)
    return jsonify({"message": "Pth initialized successfully"})

@app.route('/init_ckpt', methods=['POST'])
def init_ckpt():
    data = request.get_json()
    weights_path = data.get("path")
    if not weights_path:
        return jsonify({"error": "No weights_path provided"}), 400
    
    client = Client("http://localhost:9872/")
    result = client.predict(
        weights_path=weights_path,
        api_name="/init_t2s_weights"
    )
    print(result)
    return jsonify({"message": "Ckpt initialized successfully"})


@app.route('/api/models/list', methods=['GET'])
def change_choices():
    client = Client("http://localhost:9872/")
    result = client.predict(
            api_name="/change_choices"
    )

    try:
        message = result
        if len(message) >= 2:
            # Extract pth models from first update
            pth_models = [item[0] for item in message[0].get("choices", [])]
            # Extract ckpt models from second update
            ckpt_models = [item[0] for item in message[1].get("choices", [])]
            
            # Save to separate files
            with open("models_pth.json", "w", encoding="utf-8") as f:
                json.dump(pth_models, f, ensure_ascii=False, indent=2)
            
            with open("models_ckpt.json", "w", encoding="utf-8") as f:
                json.dump(ckpt_models, f, ensure_ascii=False, indent=2)
            
            return jsonify({
                "pth_models": pth_models,
                "ckpt_models": ckpt_models,
            })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

def signal_handler(sig, frame):
    print('Shutting down gracefully...')
    sys.exit(0)


if __name__ == '__main__':
    # Load saved configuration
    config = load_config()
    current_model = config.get("current_model")
    
    # Set up signal handler for graceful shutdown
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    # Start the Flask server
    app.run(host='0.0.0.0', port=5000, debug=True)